// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <memory>
#include <vector>

#include <nebula/compute/kernel.h>
#include <nebula/core/datum.h>

#include <turbo/base/macros.h>

namespace nebula::compute {

    /// \brief A segment
    /// A segment group is a chunk of continuous rows that have the same segment key. (For
    /// example, in ordered time series processing, segment key can be "date", and a segment
    /// group can be all the rows that belong to the same date.) A segment group can span
    /// across multiple exec batches. A segment is a chunk of continuous rows that has the
    /// same segment key within a given batch. When a segment group span cross batches, it
    /// will have multiple segments. A segment never spans cross batches. The segment data
    /// structure only makes sense when used along with a exec batch.
    struct TURBO_EXPORT Segment {
        /// \brief the offset into the batch where the segment starts
        int64_t offset;
        /// \brief the length of the segment
        int64_t length;
        /// \brief whether the segment may be extended by a next one
        bool is_open;
        /// \brief whether the segment extends a preceeding one
        bool extends;
    };

    inline bool operator==(const Segment &segment1, const Segment &segment2) {
        return segment1.offset == segment2.offset && segment1.length == segment2.length &&
               segment1.is_open == segment2.is_open && segment1.extends == segment2.extends;
    }

    inline bool operator!=(const Segment &segment1, const Segment &segment2) {
        return !(segment1 == segment2);
    }

    /// \brief a helper class to divide a batch into segments of equal values
    ///
    /// For example, given a batch with two rows:
    ///
    /// A A
    /// A A
    /// A B
    /// A B
    /// A A
    ///
    /// Then the batch could be divided into 3 segments.  The first would be rows 0 & 1,
    /// the second would be rows 2 & 3, and the third would be row 4.
    ///
    /// Further, a segmenter keeps track of the last value seen.  This allows it to calculate
    /// segments which span batches.  In our above example the last batch we emit would set
    /// the "open" flag, which indicates whether the segment may extend into the next batch.
    ///
    /// If the next call to the segmenter starts with `A A` then that segment would set the
    /// "extends" flag, which indicates whether the segment continues the last open batch.
    class TURBO_EXPORT RowSegmenter {
    public:
        virtual ~RowSegmenter() = default;

        /// \brief Construct a Segmenter which segments on the specified key types
        ///
        /// \param[in] key_types the specified key types
        /// \param[in] nullable_keys whether values of the specified keys may be null
        /// \param[in] ctx the execution context to use
        static turbo::Result<std::unique_ptr<RowSegmenter>> create(
                const std::vector<TypeHolder> &key_types, bool nullable_keys, ExecContext *ctx);

        /// \brief Return the key types of this segmenter
        virtual const std::vector<TypeHolder> &key_types() const = 0;

        /// \brief reset this segmenter
        ///
        /// A segmenter normally extends (see `Segment`) a segment from one batch to the next.
        /// If segment-extension is undesirable, for example when each batch is processed
        /// independently, then `reset` should be invoked before processing the next batch.
        virtual turbo::Status reset() = 0;

        /// \brief Get the next segment for the given batch starting from the given offset
        virtual turbo::Result<Segment> GetNextSegment(const ExecSpan &batch, int64_t offset) = 0;
    };

    /// Consumes batches of keys and yields batches of the group ids.
    class TURBO_EXPORT Grouper {
    public:
        virtual ~Grouper() = default;

        /// Construct a Grouper which receives the specified key types
        static turbo::Result<std::unique_ptr<Grouper>> create(const std::vector<TypeHolder> &key_types,
                                                              ExecContext *ctx = default_exec_context());

        /// reset all intermediate state, make the grouper logically as just `Make`ed.
        /// The underlying buffers, if any, may or may not be released though.
        virtual turbo::Status reset() = 0;

        /// consume a batch of keys, producing the corresponding group ids as an integer array,
        /// over a slice defined by an offset and length, which defaults to the batch length.
        /// Currently only uint32 indices will be produced, eventually the bit width will only
        /// be as wide as necessary.
        virtual turbo::Result<Datum> consume(const ExecSpan &batch, int64_t offset = 0,
                                             int64_t length = -1) = 0;

        /// Get current unique keys. May be called multiple times.
        virtual turbo::Result<ExecBatch> GetUniques() = 0;

        /// Get the current number of groups.
        virtual uint32_t num_groups() const = 0;

        /// \brief Assemble lists of indices of identical elements.
        ///
        /// \param[in] ids An unsigned, all-valid integral array which will be
        ///                used as grouping criteria.
        /// \param[in] num_groups An upper bound for the elements of ids
        /// \param[in] ctx Execution context to use during the operation
        /// \return A num_groups-long ListArray where the slot at i contains a
        ///         list of indices where i appears in ids.
        ///
        ///   MakeGroupings([
        ///       2,
        ///       2,
        ///       5,
        ///       5,
        ///       2,
        ///       3
        ///   ], 8) == [
        ///       [],
        ///       [],
        ///       [0, 1, 4],
        ///       [5],
        ///       [],
        ///       [2, 3],
        ///       [],
        ///       []
        ///   ]
        static turbo::Result<std::shared_ptr<ListArray>> MakeGroupings(
                const UInt32Array &ids, uint32_t num_groups,
                ExecContext *ctx = default_exec_context());

        /// \brief Produce a ListArray whose slots are selections of `array` which correspond to
        /// the provided groupings.
        ///
        /// For example,
        ///   ApplyGroupings([
        ///       [],
        ///       [],
        ///       [0, 1, 4],
        ///       [5],
        ///       [],
        ///       [2, 3],
        ///       [],
        ///       []
        ///   ], [2, 2, 5, 5, 2, 3]) == [
        ///       [],
        ///       [],
        ///       [2, 2, 2],
        ///       [3],
        ///       [],
        ///       [5, 5],
        ///       [],
        ///       []
        ///   ]
        static turbo::Result<std::shared_ptr<ListArray>> ApplyGroupings(
                const ListArray &groupings, const Array &array,
                ExecContext *ctx = default_exec_context());
    };


}  // namespace nebula::compute
